from flask import Flask, request, jsonify
from pathlib import Path
import os
from datetime import datetime
from werkzeug.utils import secure_filename
from main import CoinRecognitionSystem
from PIL import Image
import logging
import atexit

app = Flask(__name__)

# 配置上传文件的目录
UPLOAD_FOLDER = 'upload_images'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg'}

# 确保上传目录存在
Path(UPLOAD_FOLDER).mkdir(parents=True, exist_ok=True)

# 创建系统实例
system = None

def init_system():
    global system
    if system is None:
        system = CoinRecognitionSystem()
        # 加载预训练的数据库
        try:
            system.database.load_from_file('coin_database.json')
            logging.info("Successfully loaded database")
        except Exception as e:
            logging.error(f"Failed to load database: {e}")

def cleanup():
    global system
    if system and hasattr(system.database, 'client'):
        try:
            system.database.client.close()
            logging.info("Closed Milvus connection")
        except Exception as e:
            logging.error(f"Error closing Milvus connection: {e}")

# 注册清理函数
atexit.register(cleanup)

def allowed_file(filename):
    """检查文件扩展名是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/recognize', methods=['POST'])
def recognize_coin():
    """识别硬币的API端点"""
    global system
    if system is None:
        init_system()
        
    try:
        # 检查是否有正面图片
        if 'front_image' not in request.files:
            return jsonify({
                'code': 400,
                'message': '没有上传正面图片'
            }), 400
        
        # 检查是否有背面图片
        if 'back_image' not in request.files:
            return jsonify({
                'code': 400,
                'message': '没有上传背面图片'
            }), 400
        
        front_file = request.files['front_image']
        back_file = request.files['back_image']
        
        # 检查文件名是否为空
        if front_file.filename == '' or back_file.filename == '':
            return jsonify({
                'code': 400,
                'message': '没有选择文件'
            }), 400
        
        # 检查文件类型
        if not allowed_file(front_file.filename) or not allowed_file(back_file.filename):
            return jsonify({
                'code': 400,
                'message': '不支持的文件类型'
            }), 400
        
        # 生成唯一的文件名
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        timestamp_directory = Path(UPLOAD_FOLDER) / timestamp
        timestamp_directory.mkdir(parents=True, exist_ok=True)
        
        front_filepath = timestamp_directory / f"front_{secure_filename(front_file.filename)}"
        back_filepath = timestamp_directory / f"back_{secure_filename(back_file.filename)}"
        
        # 保存文件
        front_file.save(front_filepath)
        back_file.save(back_filepath)
        
        logging.info(f"Saved images to {front_filepath} and {back_filepath}")
        
        # 识别硬币
        results = system.recognize_coin(front_filepath, back_filepath)
        logging.info(f"Recognition results: {results}")
        
        if not results:
            return jsonify({
                'code': 200,
                'message': 'No similar coins found',
                'data': {
                    'results': []
                }
            })

        # 格式化返回结果
        formatted_results = []
        for result in results:
            formatted_result = {
                'coin_id': result.get('coin_id', ''),
                'combined_similarity': float(result.get('combined_similarity', 0)),
                'front_similarity': float(result.get('front_similarity', 0)) if result.get('front_similarity') is not None else 0,
                'back_similarity': float(result.get('back_similarity', 0)) if result.get('back_similarity') is not None else 0,
                'match_type': result.get('match_type', 'unknown')
            }
            formatted_results.append(formatted_result)
            
        return jsonify({
            'code': 200,
            'message': 'success',
            'data': {
                'results': formatted_results
            }
        })
        
    except Exception as e:
        logging.error(f"Error in recognize_coin: {str(e)}", exc_info=True)
        return jsonify({
            'code': 500,
            'message': f'处理失败: {str(e)}'
        }), 500

# 添加健康检查端点
@app.route('/health', methods=['GET'])
def health_check():
    return jsonify({
        'status': 'healthy',
        'timestamp': datetime.now().isoformat()
    })

if __name__ == '__main__':
    # 初始化系统
    init_system()
    
    # 关闭 Flask 的调试模式重载功能
    app.run(host='0.0.0.0', port=5001, debug=True, use_reloader=False) 